Skip to main content

media_pp\elements\filter/
audio_resampler.rs

1use std::sync::Arc;
2
3use crate::pp_log::{PpLog, pp_error, pp_info};
4use ffmpeg_next as ffmpeg;
5use thiserror::Error as ThisError;
6
7use crate::{
8    buffer::MediaBuffer,
9    control::ControlMsg,
10    element::{Element, ElementType, Sink, Source, element_pp_log},
11    error::Result,
12    pad::SrcPad,
13    time::{InvalidTimeBase, MediaTimestamp, TimeBase},
14};
15
16/// A complete uncompressed-audio format description.
17///
18/// Unlike a `(sample_rate, channels)` tuple, this also carries the sample
19/// representation and channel layout, so it can be passed directly from a
20/// hardware endpoint such as `WasapiRenderer` to an
21/// [`AudioResampler`] without guessing either one.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct AudioFormat {
24    pub sample_format: ffmpeg::format::Sample,
25    pub sample_rate: u32,
26    pub channels: u16,
27}
28
29impl AudioFormat {
30    pub fn new(sample_format: ffmpeg::format::Sample, sample_rate: u32, channels: u16) -> Self {
31        Self {
32            sample_format,
33            sample_rate,
34            channels,
35        }
36    }
37
38    pub fn channels(self) -> u16 {
39        self.channels
40    }
41
42    pub fn channel_layout(self) -> ffmpeg::ChannelLayout {
43        ffmpeg::ChannelLayout::default(i32::from(self.channels))
44    }
45}
46
47/// The reusable `libswresample` state shared by [`AudioResampler`] and
48/// [`crate::elements::SwAudioEncoder`]. It owns the otherwise easy-to-get-
49/// subtly-wrong parts of resampling: rebuilding when the input definition
50/// changes, draining the old context first, sizing output for upsampling,
51/// and flushing delayed samples at EOS.
52pub(crate) struct AudioFrameResampler {
53    target: AudioFormat,
54    context: Option<ffmpeg::software::resampling::Context>,
55}
56
57impl AudioFrameResampler {
58    pub(crate) fn new(target: AudioFormat) -> Self {
59        Self {
60            target,
61            context: None,
62        }
63    }
64
65    fn context_matches(&self, frame: &ffmpeg::frame::Audio) -> bool {
66        let Some(context) = &self.context else {
67            return false;
68        };
69        let input = context.input();
70        input.format == frame.format()
71            && input.channel_layout == frame.channel_layout()
72            && input.rate == frame.rate()
73    }
74
75    pub(crate) fn run(
76        &mut self,
77        frame: &ffmpeg::frame::Audio,
78    ) -> std::result::Result<Vec<ffmpeg::frame::Audio>, ffmpeg::Error> {
79        let mut output_frames = Vec::new();
80        if !self.context_matches(frame) {
81            output_frames.extend(self.flush()?);
82            self.context = Some(ffmpeg::software::resampling::Context::get(
83                frame.format(),
84                frame.channel_layout(),
85                frame.rate(),
86                self.target.sample_format,
87                self.target.channel_layout(),
88                self.target.sample_rate,
89            )?);
90        }
91
92        let context = self
93            .context
94            .as_mut()
95            .expect("a resampling context was built above");
96        let input_rate = frame.rate().max(1) as usize;
97        let scaled_samples = frame
98            .samples()
99            .saturating_mul(self.target.sample_rate as usize)
100            .div_ceil(input_rate);
101        let delayed_samples = context
102            .delay()
103            .map(|delay| delay.output.max(0) as usize)
104            .unwrap_or(0);
105        let capacity = scaled_samples
106            .saturating_add(delayed_samples)
107            .saturating_add(32)
108            .max(1);
109        let mut output = ffmpeg::frame::Audio::new(
110            self.target.sample_format,
111            capacity,
112            self.target.channel_layout(),
113        );
114        context.run(frame, &mut output)?;
115        if output.samples() > 0 {
116            output_frames.push(output);
117        }
118        Ok(output_frames)
119    }
120
121    pub(crate) fn flush(
122        &mut self,
123    ) -> std::result::Result<Vec<ffmpeg::frame::Audio>, ffmpeg::Error> {
124        let Some(mut context) = self.context.take() else {
125            return Ok(Vec::new());
126        };
127        let mut frames = Vec::new();
128        loop {
129            let capacity = context
130                .delay()
131                .map(|delay| delay.output.max(0) as usize)
132                .unwrap_or(0)
133                .max(1024);
134            let mut output = ffmpeg::frame::Audio::new(
135                self.target.sample_format,
136                capacity,
137                self.target.channel_layout(),
138            );
139            let delay = context.flush(&mut output)?;
140            if output.samples() > 0 {
141                frames.push(output);
142            }
143            if delay.is_none() {
144                return Ok(frames);
145            }
146        }
147    }
148
149    pub(crate) fn reset(&mut self) {
150        self.context = None;
151    }
152}
153
154/// Errors specific to [`AudioResampler`].
155#[derive(Debug, ThisError)]
156pub enum AudioResamplerError {
157    #[error("ffmpeg error: {0}")]
158    Ffmpeg(#[from] ffmpeg::Error),
159
160    #[error(
161        "AudioResampler only converts decoded Audio frames, got a {0}; link it after an audio decoder or source"
162    )]
163    UnsupportedBuffer(&'static str),
164
165    #[error(
166        "invalid input time base {numerator}/{denominator}: both numerator and denominator must be positive"
167    )]
168    InvalidTimeBase { numerator: i32, denominator: i32 },
169}
170
171/// Converts decoded audio to one fixed [`AudioFormat`] via
172/// `libswresample`. The input definition is learned from each frame and the
173/// conversion context is rebuilt if it changes mid-stream.
174///
175/// Output timestamps use `1 / target.sample_rate` units and remain
176/// contiguous across resampler buffering. The first output is anchored to
177/// the first input frame's PTS, rescaled from the explicitly supplied
178/// input time base. A decoded frame does not carry that unit itself, so
179/// callers must pass the originating stream/source time base.
180pub struct AudioResampler {
181    pp_log: PpLog,
182    name: Arc<str>,
183    target: AudioFormat,
184    input_time_base: TimeBase,
185    resampler: AudioFrameResampler,
186    next_pts: Option<i64>,
187    pad: SrcPad,
188}
189
190impl AudioResampler {
191    pub fn new(
192        name: impl Into<String>,
193        target: AudioFormat,
194        input_time_base: ffmpeg::Rational,
195    ) -> std::result::Result<Self, AudioResamplerError> {
196        let name: Arc<str> = name.into().into();
197        let pp_log = element_pp_log(ElementType::AudioResampler, &name, None);
198        let pad = SrcPad::new(format!("{name}_src"));
199        pp_info!(
200            pp_log: &pp_log,
201            "created: {}Hz, {} channel(s), format={:?}",
202            target.sample_rate,
203            target.channels(),
204            target.sample_format
205        );
206        let input_time_base = TimeBase::try_new(input_time_base).map_err(
207            |InvalidTimeBase {
208                 numerator,
209                 denominator,
210             }| AudioResamplerError::InvalidTimeBase {
211                numerator,
212                denominator,
213            },
214        )?;
215        Ok(Self {
216            name,
217            pp_log,
218            target,
219            input_time_base,
220            resampler: AudioFrameResampler::new(target),
221            next_pts: None,
222            pad,
223        })
224    }
225
226    pub fn format(&self) -> AudioFormat {
227        self.target
228    }
229
230    pub fn time_base(&self) -> ffmpeg::Rational {
231        ffmpeg::Rational::new(1, self.target.sample_rate as i32)
232    }
233
234    fn anchor_pts(&mut self, input: &ffmpeg::frame::Audio) {
235        if self.next_pts.is_some() {
236            return;
237        }
238        let pts = input
239            .pts()
240            .map(|pts| {
241                MediaTimestamp::new_unchecked(pts, self.input_time_base).rescale(
242                    TimeBase::new_unchecked(ffmpeg::Rational::new(
243                        1,
244                        self.target.sample_rate as i32,
245                    )),
246                )
247            })
248            .unwrap_or(0);
249        self.next_pts = Some(pts);
250    }
251
252    fn push_frames(&mut self, frames: Vec<ffmpeg::frame::Audio>) -> Result<()> {
253        for mut frame in frames {
254            let pts = self.next_pts.get_or_insert(0);
255            frame.set_rate(self.target.sample_rate);
256            frame.set_pts(Some(*pts));
257            *pts += frame.samples() as i64;
258            self.pad.push(MediaBuffer::Audio(Arc::new(frame)))?;
259        }
260        Ok(())
261    }
262
263    fn reset(&mut self) {
264        self.resampler.reset();
265        self.next_pts = None;
266    }
267}
268
269impl Element for AudioResampler {
270    fn name(&self) -> Arc<str> {
271        self.name.clone()
272    }
273
274    fn element_type(&self) -> ElementType {
275        ElementType::AudioResampler
276    }
277
278    fn pp_log(&self) -> &PpLog {
279        &self.pp_log
280    }
281
282    fn pp_log_mut(&mut self) -> &mut PpLog {
283        &mut self.pp_log
284    }
285}
286
287impl Source for AudioResampler {
288    fn src_pads(&mut self) -> &mut [SrcPad] {
289        std::slice::from_mut(&mut self.pad)
290    }
291}
292
293impl Sink for AudioResampler {
294    fn consume(&mut self, buf: MediaBuffer) -> Result<()> {
295        match buf {
296            MediaBuffer::Audio(frame) => {
297                self.anchor_pts(&frame);
298                let frames = self
299                    .resampler
300                    .run(&frame)
301                    .inspect_err(|error| pp_error!(self, "resample failed: {error}"))
302                    .map_err(AudioResamplerError::from)?;
303                self.push_frames(frames)
304            }
305            MediaBuffer::Eos => {
306                let frames = self
307                    .resampler
308                    .flush()
309                    .inspect_err(|error| pp_error!(self, "resampler flush failed: {error}"))
310                    .map_err(AudioResamplerError::from)?;
311                self.push_frames(frames)?;
312                self.pad.push(MediaBuffer::Eos)
313            }
314            MediaBuffer::Packet(_) => Err(AudioResamplerError::UnsupportedBuffer("Packet").into()),
315            MediaBuffer::Video(_) => Err(AudioResamplerError::UnsupportedBuffer("Video").into()),
316        }
317    }
318
319    fn control(&mut self, msg: ControlMsg) -> Result<()> {
320        if matches!(msg, ControlMsg::Seek(_) | ControlMsg::Stop) {
321            self.reset();
322        }
323        self.pad.control(msg)
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use std::sync::Mutex;
330
331    use ffmpeg::format::sample::Type;
332
333    use super::*;
334
335    struct CapturingSink {
336        pp_log: PpLog,
337        received: Arc<Mutex<Vec<MediaBuffer>>>,
338    }
339
340    impl Element for CapturingSink {
341        fn name(&self) -> Arc<str> {
342            "capture".into()
343        }
344
345        fn element_type(&self) -> ElementType {
346            ElementType::Other
347        }
348
349        fn pp_log(&self) -> &PpLog {
350            &self.pp_log
351        }
352
353        fn pp_log_mut(&mut self) -> &mut PpLog {
354            &mut self.pp_log
355        }
356    }
357
358    impl Sink for CapturingSink {
359        fn consume(&mut self, buf: MediaBuffer) -> Result<()> {
360            self.received.lock().unwrap().push(buf);
361            Ok(())
362        }
363
364        fn control(&mut self, _msg: ControlMsg) -> Result<()> {
365            Ok(())
366        }
367    }
368
369    fn f32_packed_frame(rate: u32, channels: u16, samples: usize, pts: i64) -> MediaBuffer {
370        let format = ffmpeg::format::Sample::F32(Type::Packed);
371        let layout = ffmpeg::ChannelLayout::default(i32::from(channels));
372        let mut frame = ffmpeg::frame::Audio::new(format, samples, layout);
373        frame.set_rate(rate);
374        frame.set_pts(Some(pts));
375        frame.data_mut(0).fill(0);
376        MediaBuffer::Audio(Arc::new(frame))
377    }
378
379    fn new_resampler(target: AudioFormat) -> (AudioResampler, Arc<Mutex<Vec<MediaBuffer>>>) {
380        let mut resampler =
381            AudioResampler::new("resampler", target, ffmpeg::Rational::new(1, 48_000)).unwrap();
382        let received = Arc::new(Mutex::new(Vec::new()));
383        resampler.src_pads()[0].link(Box::new(CapturingSink {
384            received: received.clone(),
385            pp_log: element_pp_log(ElementType::Other, "capture", None),
386        }));
387        (resampler, received)
388    }
389
390    #[test]
391    fn converts_rate_channels_and_timestamps_then_flushes_eos() {
392        let target = AudioFormat::new(ffmpeg::format::Sample::F32(Type::Packed), 24_000, 1);
393        let (mut resampler, received) = new_resampler(target);
394        resampler
395            .consume(f32_packed_frame(48_000, 2, 960, 48_000))
396            .unwrap();
397        resampler.consume(MediaBuffer::Eos).unwrap();
398
399        let received = received.lock().unwrap();
400        let audio: Vec<_> = received
401            .iter()
402            .filter_map(|buffer| match buffer {
403                MediaBuffer::Audio(frame) => Some(frame),
404                _ => None,
405            })
406            .collect();
407        assert!(!audio.is_empty());
408        assert_eq!(audio[0].pts(), Some(24_000));
409        assert!(
410            audio
411                .iter()
412                .all(|frame| frame.format() == target.sample_format)
413        );
414        assert!(audio.iter().all(|frame| frame.rate() == 24_000));
415        assert!(audio.iter().all(|frame| frame.channels() == 1));
416        for pair in audio.windows(2) {
417            assert_eq!(
418                pair[1].pts(),
419                pair[0].pts().map(|pts| pts + pair[0].samples() as i64)
420            );
421        }
422        assert!(matches!(received.last(), Some(MediaBuffer::Eos)));
423    }
424
425    #[test]
426    fn rebuilds_when_the_input_definition_changes() {
427        let target = AudioFormat::new(ffmpeg::format::Sample::F32(Type::Packed), 48_000, 2);
428        let (mut resampler, received) = new_resampler(target);
429        resampler
430            .consume(f32_packed_frame(48_000, 2, 480, 0))
431            .unwrap();
432        resampler
433            .consume(f32_packed_frame(44_100, 1, 441, 441))
434            .unwrap();
435        resampler.consume(MediaBuffer::Eos).unwrap();
436
437        assert!(received.lock().unwrap().iter().any(|buffer| {
438            matches!(buffer, MediaBuffer::Audio(frame) if frame.rate() == 48_000 && frame.channels() == 2)
439        }));
440    }
441
442    #[test]
443    fn rejects_non_audio_buffers() {
444        let target = AudioFormat::new(ffmpeg::format::Sample::F32(Type::Packed), 48_000, 2);
445        let (mut resampler, _) = new_resampler(target);
446        let error = resampler
447            .consume(MediaBuffer::Packet(Arc::new(ffmpeg::Packet::empty())))
448            .unwrap_err();
449        assert!(matches!(
450            error,
451            crate::Error::AudioResamplerError(AudioResamplerError::UnsupportedBuffer("Packet"))
452        ));
453    }
454}